// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/expression/decoded_args.h>
#include <pollux/expression/eval_ctx.h>
#include <pollux/expression/vector_function.h>

namespace kumo::pollux::functions {
    namespace {
        template<bool IsNotNULL>
        class IsNullFunction : public exec::VectorFunction {
        public:
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr & /*outputType*/,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                auto *arg = args[0].get();
                auto *pool = context.pool();
                if (arg->is_constant_encoding()) {
                    bool isNull = arg->is_null_at(rows.begin());
                    auto localResult = BaseVector::create_constant(
                        BOOLEAN(), IsNotNULL ? !isNull : isNull, rows.end(), pool);
                    context.moveOrCopyResult(localResult, rows, result);
                    return;
                }

                if (!arg->may_have_nulls()) {
                    // No nulls.
                    auto localResult = BaseVector::create_constant(
                        BOOLEAN(), IsNotNULL ? true : false, rows.end(), pool);
                    context.moveOrCopyResult(localResult, rows, result);
                    return;
                }

                BufferPtr isNull;
                if (arg->is_flat_encoding()) {
                    if constexpr (IsNotNULL) {
                        isNull = arg->nulls();
                    } else {
                        isNull = AlignedBuffer::allocate<bool>(rows.end(), pool);
                        memcpy(
                            isNull->asMutable<int64_t>(),
                            arg->raw_nulls(),
                            bits::nbytes(rows.end()));
                        bits::negate(isNull->asMutable<uint64_t>(), rows.end());
                    }
                } else {
                    exec::DecodedArgs decodedArgs(rows, args, context);

                    isNull = AlignedBuffer::allocate<bool>(rows.end(), pool);
                    memcpy(
                        isNull->asMutable<int64_t>(),
                        decodedArgs.at(0)->nulls(&rows),
                        bits::nbytes(rows.end()));

                    if (!IsNotNULL) {
                        bits::negate(isNull->asMutable<uint64_t>(), rows.end());
                    }
                }

                auto localResult = std::make_shared<FlatVector<bool> >(
                    pool, BOOLEAN(), nullptr, rows.end(), isNull, std::vector<BufferPtr>{});
                context.moveOrCopyResult(localResult, rows, result);
            }

            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                // T -> boolean
                return {
                    exec::FunctionSignatureBuilder()
                    .typeVariable("T")
                    .returnType("boolean")
                    .argumentType("T")
                    .build()
                };
            }
        };
    } // namespace

    POLLUX_DECLARE_VECTOR_FUNCTION_WITH_METADATA(
        udf_is_null,
        IsNullFunction<false>::signatures(),
        exec::VectorFunctionMetadataBuilder().defaultNullBehavior(false).build(),
        std::make_unique<IsNullFunction</*IsNotNUll=*/false>>());

    void registerIsNullFunction(const std::string &name) {
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_is_null, name);
    }

    POLLUX_DECLARE_VECTOR_FUNCTION_WITH_METADATA(
        udf_is_not_null,
        IsNullFunction<true>::signatures(),
        exec::VectorFunctionMetadataBuilder().defaultNullBehavior(false).build(),
        std::make_unique<IsNullFunction</*IsNotNUll=*/true>>());

    void registerIsNotNullFunction(const std::string &name) {
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_is_not_null, name);
    }
} // namespace kumo::pollux::functions
